// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/connectors/tpch/tpch_connector.h>
#include <pollux/tpch/gen/tpch_gen.h>

namespace kumo::pollux::connector::tpch {

using kumo::pollux::tpch::Table;

namespace {

RowVectorPtr getTpchData(
    Table table,
    size_t maxRows,
    size_t offset,
    double scaleFactor,
    memory::MemoryPool* pool) {
  switch (table) {
    case Table::TBL_PART:
      return pollux::tpch::genTpchPart(pool, maxRows, offset, scaleFactor);
    case Table::TBL_SUPPLIER:
      return pollux::tpch::genTpchSupplier(pool, maxRows, offset, scaleFactor);
    case Table::TBL_PARTSUPP:
      return pollux::tpch::genTpchPartSupp(pool, maxRows, offset, scaleFactor);
    case Table::TBL_CUSTOMER:
      return pollux::tpch::genTpchCustomer(pool, maxRows, offset, scaleFactor);
    case Table::TBL_ORDERS:
      return pollux::tpch::genTpchOrders(pool, maxRows, offset, scaleFactor);
    case Table::TBL_LINEITEM:
      return pollux::tpch::genTpchLineItem(pool, maxRows, offset, scaleFactor);
    case Table::TBL_NATION:
      return pollux::tpch::genTpchNation(pool, maxRows, offset, scaleFactor);
    case Table::TBL_REGION:
      return pollux::tpch::genTpchRegion(pool, maxRows, offset, scaleFactor);
  }
  return nullptr;
}

} // namespace

std::string TpchTableHandle::toString() const {
  return fmt::format(
      "table: {}, scale factor: {}", toTableName(table_), scaleFactor_);
}

TpchDataSource::TpchDataSource(
    const std::shared_ptr<const RowType>& outputType,
    const std::shared_ptr<connector::ConnectorTableHandle>& tableHandle,
    const std::unordered_map<
        std::string,
        std::shared_ptr<connector::ColumnHandle>>& columnHandles,
    pollux::memory::MemoryPool* pool)
    : pool_(pool) {
  auto tpchTableHandle =
      std::dynamic_pointer_cast<TpchTableHandle>(tableHandle);
  POLLUX_CHECK_NOT_NULL(
      tpchTableHandle, "TableHandle must be an instance of TpchTableHandle");
  tpchTable_ = tpchTableHandle->getTable();
  scaleFactor_ = tpchTableHandle->getScaleFactor();
  tpchTableRowCount_ = getRowCount(tpchTable_, scaleFactor_);

  auto tpchTableSchema = getTableSchema(tpchTableHandle->getTable());
  POLLUX_CHECK_NOT_NULL(tpchTableSchema, "TpchSchema can't be null.");

  outputColumnMappings_.reserve(outputType->size());

  for (const auto& outputName : outputType->names()) {
    auto it = columnHandles.find(outputName);
    POLLUX_CHECK(
        it != columnHandles.end(),
        "ColumnHandle is missing for output column '{}' on table '{}'",
        outputName,
        toTableName(tpchTable_));

    auto handle = std::dynamic_pointer_cast<TpchColumnHandle>(it->second);
    POLLUX_CHECK_NOT_NULL(
        handle,
        "ColumnHandle must be an instance of TpchColumnHandle "
        "for '{}' on table '{}'",
        it->second->name(),
        toTableName(tpchTable_));

    auto idx = tpchTableSchema->getChildIdxIfExists(handle->name());
    POLLUX_CHECK(
        idx != std::nullopt,
        "Column '{}' not found on TPC-H table '{}'.",
        handle->name(),
        toTableName(tpchTable_));
    outputColumnMappings_.emplace_back(*idx);
  }
  outputType_ = outputType;
}

RowVectorPtr TpchDataSource::projectOutputColumns(RowVectorPtr inputVector) {
  std::vector<VectorPtr> children;
  children.reserve(outputColumnMappings_.size());

  for (const auto channel : outputColumnMappings_) {
    children.emplace_back(inputVector->childAt(channel));
  }

  return std::make_shared<RowVector>(
      pool_,
      outputType_,
      BufferPtr(),
      inputVector->size(),
      std::move(children));
}

void TpchDataSource::addSplit(std::shared_ptr<ConnectorSplit> split) {
  POLLUX_CHECK_EQ(
      currentSplit_,
      nullptr,
      "Previous split has not been processed yet. Call next() to process the split.");
  currentSplit_ = std::dynamic_pointer_cast<TpchConnectorSplit>(split);
  POLLUX_CHECK(currentSplit_, "Wrong type of split for TpchDataSource.");

  // Lineitems is generated based on the row ids of the orders table, so
  // splitOffset_ and splitEnd_ will refer to orders.
  size_t effectiveRowCount = isLineItem()
      ? getRowCount(Table::TBL_ORDERS, scaleFactor_)
      : tpchTableRowCount_;

  size_t partSize = std::ceil(
      static_cast<double>(effectiveRowCount) /
      static_cast<double>(currentSplit_->totalParts));

  splitOffset_ = partSize * currentSplit_->partNumber;
  splitEnd_ = splitOffset_ + partSize;
}

std::optional<RowVectorPtr> TpchDataSource::next(
    uint64_t size,
    pollux::ContinueFuture& /*future*/) {
  POLLUX_CHECK_NOT_NULL(
      currentSplit_, "No split to process. Call addSplit() first.");

  // LineItems generates records based on orders, so it will generate on
  // average 4 times more records than what is requested. Dividing by 4 so it
  // generates about the right amount of records. Note that the exact amount of
  // lineitems for an order is random (from 1 to 7), so this function may
  // generate slightly more or fewer records than `size`.
  if (isLineItem()) {
    size /= 4;
  }

  size_t maxRows = std::min(size, (splitEnd_ - splitOffset_));
  auto outputVector =
      getTpchData(tpchTable_, maxRows, splitOffset_, scaleFactor_, pool_);

  // If the split is exhausted.
  if (!outputVector || outputVector->size() == 0) {
    currentSplit_ = nullptr;
    return nullptr;
  }

  // splitOffset needs to advance based on maxRows passed to getTpchData(), and
  // not the actual number of returned rows in the output vector, as they are
  // not the same for lineitem.
  splitOffset_ += maxRows;
  completedRows_ += outputVector->size();
  completedBytes_ += outputVector->retained_size();

  return projectOutputColumns(outputVector);
}

bool TpchDataSource::isLineItem() const {
  return tpchTable_ == Table::TBL_LINEITEM;
}

} // namespace kumo::pollux::connector::tpch
